#PBAR_ExamineBaseRecovery.py

import PBAR_FD
from scipy import ndimage
from scipy.optimize import curve_fit

plotColors = ['r', 'b', 'g', 'm', 'c', 'y', 'k'] * 10
lineStyles = ['-', ':', '-.', '_', '|'] *  10
markerTypes = ['.', 'o', 'v', '^', '<', '>', '1', '2', '3', '4', 's', 'p', '*', 'h', 'H', '+', 'x', 'D', 'd']

# READ IN THE DATA
filePrefixList = ['dz33','dz34','dz35','dz36','dz37','dz38','dz39']
filePath = 'C:\Users\jkwong\Documents\Work\PBAR\data'

dat = []
for i in xrange(len(filePrefixList)):
    dat.append(PBAR_FD.ReadData(filePrefixList[i], filePath))

# Data set groups

# Get count from hydrogen peak and iron peak
# skip the first time bin

def gauss_function(x, a, x0, sigma):
    return a*np.exp(-(x-x0)**2/(2*sigma**2))

# Time binedges

# for H
timeBoundariesList1 = np.array((2, 3, 4, 5, 6, 7, 9, 11))
# for Fe
timeBoundariesList2 = np.hstack((np.arange(2,20,2), np.arange(20,60,10)))

datasetIndices = [6, 6, 6, 6]
detChList = [3, 3, 28, 28]
timeBinList = [timeBoundariesList1, timeBoundariesList2, timeBoundariesList1, timeBoundariesList2]

# location of the peak
peakLocation1 = [42.0, 42.0, 42.0, 42.0, 42.0, 42.0, 42.0, 42.0]
peakLocation2 = [145.6, 148.6, 149.5, 148.3,150.69, 150.40, 148.9, 148.30,146.80, 147.4,139.9, 138.7]

peakBounds1 = [[35,45], [35,44], [35,44], [35,44], [35,44], [35,44], [35,44], [35,44]]
#peakBounds2 = [[134,154],[134,154],[134,154],[134,154],[134,154],[132,154],[132,154],[132,154],[132,154],[131,154],[131,154],[131,154],]
peakBounds2 = [[122,154],[122,154],[122,154],[122,154],[122,154],[122,154],[122,154],[122,154],[122,154],[122,154],[122,154],[122,154],]

peakBoundsList = [peakBounds1, peakBounds2, peakBounds1, peakBounds2]
peakLocationList = [peakLocation1, peakLocation2, peakLocation1, peakLocation2]

#subtractBackground = [1, 0, 1, 0]
backgroundRegion = [60, 80]

fitParametersList = []
backgroundList = []
backgroundpfitList = []

for j in xrange(len(datasetIndices)):
    
    index = datasetIndices[j]
    detNum = detChList[j]
    timeBoundariesList = timeBinList[j]
    peakBounds = peakBoundsList[j]
    peakLocation = peakLocationList[j]
    
    # make temporary containers
    fitParameters = []
    background = []
    backgroundpfit = []
    
    for i in xrange(len(timeBoundariesList) - 1):
        y = dat[index][detNum][:,timeBoundariesList[i]:timeBoundariesList[i+1]].sum(1)
        y = y / (timeBoundariesList[i+1] - timeBoundariesList[i])
        
        # Set starting parameters
        x = np.arange(256)
        cut = (x > peakBounds[i][0]) & (x < peakBounds[i][1])
        startingParam = [max(y[cut]), peakLocation[i], (peakBounds[i][1] - peakBounds[i][0])/2]
        try:
            popt, pcov = curve_fit(gauss_function, x[cut], y[cut], p0 = startingParam)
        except:
            print "Bad Fit"
            popt = [-1, -1, -1]
        
        print popt
        fitMean = popt[1]
        fitSigma = popt[2]
        fitAmp = popt[0]
        
        ## fit a second time
        #upperBound = fitMean + 2.0 * fitSigma
        #lowerBound = fitMean - 1.0 * fitSigma
        
        #cut = (x > lowerBound) & (x < upperBound)
        #startingParam = [fitAmp, fitMean, fitSigma]
        #popt, pcov = curve_fit(gauss_function, x[cut], y[cut], p0 = startingParam)
        #print popt
        
        fitParameters.append(popt)
        fitMean = popt[1]
        fitSigma = popt[2]
        fitAmp = popt[0]        
        
        # Fit to background section
        cut = (x > backgroundRegion[0]) & (x < backgroundRegion[1]) | ((x >= 17) & (x <= 19))
        cut = (x > backgroundRegion[0]) & (x < backgroundRegion[1])
        
        pfit = np.polyfit(x[cut], log(y[cut]), 1)
        background.append(exp(np.polyval(pfit, fitMean)))
        backgroundpfit.append(pfit)
    
    fitParameters = np.array(fitParameters)
    backgroundpfit = np.array(backgroundpfit)
    background = np.array(background)
    
    fitParametersList.append(fitParameters)
    backgroundList.append(background)
    backgroundpfitList.append(backgroundpfit)
    


# PLOTS



# PLOTS

# Plot the waveforms at multiple times bins, separate plots
index = 6
detNum = 28

filterSpectra = 0
filterWidth = 1

#timeBoundariesList = np.hstack((np.arange(2,20,1), np.arange(20,100,10)))
timeBoundariesList = np.hstack((np.arange(2,20,2), np.arange(20,100,10)))

detNumList = np.arange(9)

for detNum in detNumList:
    plt.figure()
    plt.grid()
    for i in xrange(len(timeBoundariesList)-1):    
        if filterSpectra:
            y = ndimage.filters.gaussian_filter(dat[index][detNum][:,timeBoundariesList[i]:timeBoundariesList[i+1]].sum(1), filterWidth, order = 0)
        else:
            y = dat[index][detNum][:,timeBoundariesList[i]:timeBoundariesList[i+1]].sum(1)
        y = y / (timeBoundariesList[i+1] - timeBoundariesList[i])
        
        plot(y, label = 't bin (%d,%d)' %(timeBoundariesList[i],timeBoundariesList[i+1]))
    
    plt.xlim((0, 255))
    plt.yscale('log')
    plt.legend()
    plt.ylabel('Counts')
    plt.xlabel('Energy')
    plt.yscale('log')
    plt.title('Dataset: %s, Detector # %d' %(filePrefixList[index], detNum+1))



params = {'legend.fontsize': 8,
          'legend.linewidth': 2}
plt.rcParams.update(params)




# Plot the waveforms at multiple times bins, subplots
index = 0
index = 6

filterSpectra = 0
filterWidth = 1
normalize = 0

timeBoundariesList = np.hstack((np.arange(2,20,2), np.arange(20,100,10)))
#timeBoundariesList = timeBoundariesList[0::2]

detNumList = np.arange(9)

subplotx = 4
subploty = 2

# cycle through detectors
for detNum in np.arange(40):
    # Generate a new figure window if filled.
    if ((detNum % (subplotx * subploty)) == 0):
        f, ax = plt.subplots(subploty, subplotx, sharex='col', sharey='row')
    # Calculate subplot window index
    ii = (detNum % (subplotx*subploty)) /subplotx
    jj = (detNum % (subplotx*subploty)) % subplotx
    
    for i in xrange(len(timeBoundariesList)-1):    
        if filterSpectra:
            y = ndimage.filters.gaussian_filter(dat[index][detNum][:,timeBoundariesList[i]:timeBoundariesList[i+1]].sum(1), filterWidth, order = 0)
        else:
            y = dat[index][detNum][:,timeBoundariesList[i]:timeBoundariesList[i+1]].sum(1)
        y = y / (timeBoundariesList[i+1] - timeBoundariesList[i])
        
        if normalize:
            y = y / sum(y[0:255])

        ax[ii][jj].plot(y, label = 't bin (%d,%d)' %(timeBoundariesList[i],timeBoundariesList[i+1]))


    ax[ii][jj].grid()
    if (ii == 0) & (jj == 0):
        ax[ii][jj].legend()
    ax[ii][jj].set_xlim((0, 255))
    #ax[ii][ii].set_ylim((1e-1, 1e5))
    ax[ii][jj].set_yscale('log')
    ax[ii][jj].set_ylabel('Counts')
    ax[ii][jj].set_xlabel('Energy')
    ax[ii][jj].set_title('Dataset: %s, Detector # %d' %(filePrefixList[index], detNum+1))




# PLOT SPECTRA - FIRST SET
timeBoundariesList = timeBoundariesList1
#timeBoundariesList = np.hstack((np.arange(2,20,2), np.arange(20,60,10)))

setIndex = 

datasetIndices = [6, 6, 6, 6]
detChList = [3, 3, 28, 28]
timeBinList = [timeBoundariesList1, timeBoundariesList2, timeBoundariesList1, timeBoundariesList2]


index = 6
detNum = 28
detNum = 3

k = 0

filterSpectra = 0
filterWidth = 1
plt.figure()
plt.grid()


i = 0
timeBoundary = [0, 1]
if filterSpectra:
    y = ndimage.filters.gaussian_filter(dat[index][detNum][:,timeBoundary[i]:timeBoundary[i+1]].sum(1), filterWidth, order = 0)
else:
    y = dat[index][detNum][:,timeBoundary[i]:timeBoundary[i+1]].sum(1)
y = y / (timeBoundary[i+1] - timeBoundary[i])

plot(y, label = 't bin (%d,%d)' %(timeBoundary[i],timeBoundary[i+1]))

timeBoundary = [1, 2]
i = 0
if filterSpectra:
    y = ndimage.filters.gaussian_filter(dat[index][detNum][:,timeBoundary[i]:timeBoundary[i+1]].sum(1), filterWidth, order = 0)
else:
    y = dat[index][detNum][:,timeBoundary[i]:timeBoundary[i+1]].sum(1)
y = y / (timeBoundary[i+1] - timeBoundary[i])

plot(y, label = 't bin (%d,%d)' %(timeBoundary[i],timeBoundary[i+1]))


for i in xrange(len(timeBoundariesList)-1):    

    if filterSpectra:
        y = ndimage.filters.gaussian_filter(dat[index][detNum][:,timeBoundariesList[i]:timeBoundariesList[i+1]].sum(1), filterWidth, order = 0)
    else:
        y = dat[index][detNum][:,timeBoundariesList[i]:timeBoundariesList[i+1]].sum(1)
    y = y / (timeBoundariesList[i+1] - timeBoundariesList[i])
    
    plot(y, label = 't bin (%d,%d)' %(timeBoundariesList[i],timeBoundariesList[i+1]))

    # over lay best fit to peak
    fitParam = fitParametersList[k]
    lowerBound = fitParam[i][1] - fitParam[i][2] * 2.0
    upperBound = fitParam[i][1] + fitParam[i][2] * 2.0
    
    xx = np.linspace(lowerBound, upperBound, 50)
    yy = gauss_function(xx, fitParam[i][0], fitParam[i][1], fitParam[i][2])

    plot(xx, yy)
    # plot the background level
    xx = np.linspace(20, 80, 50)
    pfit = backgroundpfitList[k][i,:]
    yy = exp(np.polyval(pfit, xx))
    plot(xx, yy)
    
plt.yscale('log')
plt.legend()
plt.ylabel('Counts')
plt.xlabel('Energy')
plt.yscale('log')
plt.title('Dataset: %s, Detector # %d' %(filePrefixList[index], detNum+1))






# H Peak count versus time
timeCal = 16.384e-6
k = 0
timeCenters = (timeBinList[k][0:-1] + timeBinList[k][1:])/2.0
counts = fitParametersList[k][:,0] - backgroundList[k]

# fit to exponential
pfit = np.polyfit(timeCenters*timeCal, log(counts), 1)
t = -1/pfit[0]

plt.figure();
plt.grid()
plt.plot(timeCenters*timeCal, counts, 'xk', markersize = 15, linewidth = 2.)
xx = np.linspace(0, 0.00018, 50)
yy = exp(np.polyval(pfit, xx))
plot(xx, yy)

plt.xlabel('Time (sec)')
plt.ylabel('Counts')
plt.title('H Line Decay Time = %f sec' %t)




# H Peak location versus time

timeCal = 16.384e-6

k = 0
timeCenters = (timeBinList[k][0:-1] + timeBinList[k][1:])/2.0
fitMeans = fitParametersList[k][:,1]

plt.figure();
plt.grid()
plt.plot(timeCenters*timeCal, fitMeans, 'xk', markersize = 15)

plt.xlabel('Time (sec)')
plt.ylabel('Fit Mean')
plt.title('H, 2.224 MeV')



# PLOT SPECTRA - SECOND SET - Fe peak

#timeBoundariesList = np.arange(2,11,1)
timeBoundariesList = np.hstack((np.arange(2,20,2), np.arange(20,60,10)))
index = 6
detNum = 28
detNum = 3

k = 1

filterSpectra = 0
filterWidth = 1
figure()
plt.grid()


timeBoundary = [0, 1]
i = 0
if filterSpectra:
    y = ndimage.filters.gaussian_filter(dat[index][detNum][:,timeBoundary[i]:timeBoundary[i+1]].sum(1), filterWidth, order = 0)
else:
    y = dat[index][detNum][:,timeBoundary[i]:timeBoundary[i+1]].sum(1)
y = y / (timeBoundary[i+1] - timeBoundary[i])

plot(y, label = 't bin (%d,%d)' %(timeBoundary[i],timeBoundary[i+1]))


timeBoundary = [1, 2]
i = 0
if filterSpectra:
    y = ndimage.filters.gaussian_filter(dat[index][detNum][:,timeBoundary[i]:timeBoundary[i+1]].sum(1), filterWidth, order = 0)
else:
    y = dat[index][detNum][:,timeBoundary[i]:timeBoundary[i+1]].sum(1)
y = y / (timeBoundary[i+1] - timeBoundary[i])

plot(y, label = 't bin (%d,%d)' %(timeBoundary[i],timeBoundary[i+1]))


for i in xrange(len(timeBoundariesList)-1):
    if filterSpectra:
        y = ndimage.filters.gaussian_filter(dat[index][detNum][:,timeBoundariesList[i]:timeBoundariesList[i+1]].sum(1), filterWidth, order = 0)
    else:
        y = dat[index][detNum][:,timeBoundariesList[i]:timeBoundariesList[i+1]].sum(1)
    y = y / (timeBoundariesList[i+1] - timeBoundariesList[i])
    
    # over lay best fit
    fitParam = fitParametersList[k]
    lowerBound = fitParam[i][1] - fitParam[i][2] * 2.0
    upperBound = fitParam[i][1] + fitParam[i][2] * 2.0
    
    xx = np.linspace(lowerBound, upperBound, 50)
    yy = gauss_function(xx, fitParam[i][0], fitParam[i][1], fitParam[i][2])
    
    plot(y, label = 't bin (%d,%d)' %(timeBoundariesList[i],timeBoundariesList[i+1]), color = plotColors[i], linestyle = lineStyles[i/7])
    plot(xx, yy, color = plotColors[i], linewidth = 2.0, linestyle = lineStyles[i/7])
plt.yscale('log')
plt.legend()
plt.ylabel('Counts')
plt.xlabel('Bin')
plt.yscale('log')
plt.title('Dataset: %s, Detector # %d' %(filePrefixList[index], detNum+1))



# Fe peak versus time
timeCal = 16.384e-6

k = 1
timeCenters = (timeBinList[k][0:-1] + timeBinList[k][1:])/2.0
fitMeans = fitParametersList[k][:,1]

plt.figure();
plt.grid()
plt.plot(timeCenters*timeCal, fitMeans, 'xk', markersize = 15)


plt.xlabel('Time (sec)')
plt.ylabel('Fit Mean')
plt.title('Fe, 7.645 MeV')

